#!/bin/bash

# 检查系统类型
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$ID
else
    echo "无法检测系统类型。"
    exit 1
fi

# 根据系统类型安装NVIDIA驱动程序
case $OS in
    ubuntu | debian)
        echo "检测到系统为Ubuntu或Debian。开始安装NVIDIA驱动程序。"
        # 确保没有待处理的更新：
        sudo apt-get update && sudo apt-get upgrade -y
        sleep 1
        clear

        # 添加ubuntu-drivers的仓库并选择推荐的驱动程序：
        sudo apt-get install -y ubuntu-drivers-common

        # 列出此系统可用的驱动程序：
        sudo ubuntu-drivers devices

        echo ""
        echo "请查看上面的驱动程序列表并选择您希望为系统安装的驱动程序。"
        echo "如果您希望自动选择最新的驱动程序，请直接按回车键。"
        read -p "输入您的选择: " option

        case $option in
            nvidia-driver-*)
                echo "现在使用请求的更改更新：${option}"
                sudo apt install -y ${option}
                ;;
            *)
                echo "自动选择最新的驱动程序"
                sudo ubuntu-drivers autoinstall
                ;;
        esac

        # 黑名单Nouveau驱动程序，以防止其初始化：
        sudo bash -c "echo 'blacklist nouveau' > /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
        sudo bash -c "echo 'options nouveau modeset=0' >> /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
        sleep 1
        cat /etc/modprobe.d/blacklist-nvidia-nouveau.conf
        sleep 1
        clear

        # 更新内核以反映更改：
        echo "更新initramfs..."
        sleep 1
        sudo update-initramfs -u
        clear

        echo "脚本完成 - NVIDIA驱动程序已安装"
        echo "重新启动电脑，输入：nvidia-smi 查询驱动是否安装成功"
        ;;
    centos | rhel)
        echo "检测到系统为CentOS或Red Hat。开始安装NVIDIA驱动程序。"
        # 确保没有待处理的更新：
        sudo yum update -y
        sleep 1
        clear

        # 安装EPEL仓库
        sudo yum install -y epel-release
        sudo yum update -y

        # 安装NVIDIA驱动程序仓库
        sudo yum install -y https://dl.fedoraproject.org/pub/epel/epel-release-latest-$(rpm -E %rhel).noarch.rpm
        sudo yum-config-manager --add-repo=http://developer.download.nvidia.com/compute/cuda/repos/rhel$(rpm -E %rhel)/x86_64/cuda-rhel$(rpm -E %rhel).repo

        # 安装NVIDIA驱动程序
        sudo yum clean all
        sudo yum -y install nvidia-driver-latest-dkms
        sudo yum -y install cuda

        # 禁用Nouveau驱动程序
        sudo bash -c "echo 'blacklist nouveau' > /etc/modprobe.d/blacklist-nouveau.conf"
        sudo bash -c "echo 'options nouveau modeset=0' >> /etc/modprobe.d/blacklist-nouveau.conf"
        sleep 1
        cat /etc/modprobe.d/blacklist-nouveau.conf
        sleep 1
        clear

        # 生成新的initramfs
        echo "生成新的initramfs..."
        sleep 1
        sudo dracut --force
        clear

        echo "脚本完成 - NVIDIA驱动程序已安装"
        echo "重新启动电脑，输入：nvidia-smi 查询驱动是否安装成功"
        ;;
    *)
        echo "不支持的操作系统：$OS。停止安装。"
        exit 1
        ;;
esac